import argparse
import logging
import numpy as np
import cv2
import torch
from PIL import Image
import carla
import pygame
import random

from data_generation.data_generator import DataGenerator
from data_generation.keyboard_control import KeyboardControl
from  data_generation.world import World
from data_generation.bev_render import BevRender

def world_init(args):
    client = carla.Client(args.host, args.port)
    client.set_timeout(5.0)
    world = client.get_world()
    carla_world = client.load_world(args.map)
    print('world set done')

    # 生成汽车，设置自动驾驶
    blueprint_library = world.get_blueprint_library()
    ego_vehicle_bp = blueprint_library.find('vehicle.mercedes-benz.coupe')
    ego_vehicle_bp.set_attribute('color', '0, 0, 0')
    transform = random.choice(world.get_map().get_spawn_points())
    ego_vehicle = world.spawn_actor(ego_vehicle_bp, transform)
    ego_vehicle.set_autopilot(True)

    # 设置同步模式
    settings = world.get_settings()
    settings.fixed_delta_seconds = 0.05  # 20 fps, 5ms
    world.apply_settings(settings)

    _world = World(carla_world, args)
    _world._player = ego_vehicle
    _world._bev_render = BevRender(_world, args.bev_render_device)

    return _world, ego_vehicle

def generate_bev(args, world_, ego_vehicle):
    bev_state = world_.bev_state
    bev_view = world_.render_BEV_from_state(bev_state)
    return bev_view

def str2bool(v):
    if v.lower() in ('yes', 'true', 'True', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'False', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Unsupported value encountered.')


def encode_npy_to_pil(bev_array):
    c, w, h = bev_array.shape

    img = np.zeros([3, w, h]).astype('uint8')
    bev = np.ceil(bev_array).astype('uint8')

    for i in range(c):
        if 0 <= i <= 4:
            # road, lane, light
            img[0] = img[0] | (bev[i] << (8 - i - 1))
        elif 5 <= i <= 9:
            img[1] = img[1] | (bev[i] << (8 - (i - 5) - 1))
        elif 10 <= i <= 14:
            img[2] = img[2] | (bev[i] << (8 - (i - 10) - 1))

    return img


def main():
    argparser = argparse.ArgumentParser(
        description='CARLA Data Generation')
    argparser.add_argument(
        '-v', '--verbose',
        action='store_true',
        dest='debug',
        help='print debug information')
    argparser.add_argument(
        '--host',
        metavar='H',
        default='127.0.0.1',
        help='IP of the host server (default: 127.0.0.1)')
    argparser.add_argument(
        '-p', '--port',
        metavar='P',
        default=2000,
        type=int,
        help='TCP port to listen to (default: 2000)')
    argparser.add_argument(
        '--res',
        metavar='WIDTHxHEIGHT',
        default='860x480',
        help='window resolution (default: 860x480)')
    argparser.add_argument(
        '--gamma',
        default=0.0,
        type=float,
        help='Gamma correction of the camera (default: 0.0)')
    argparser.add_argument(
        '--save_path',
        default='./e2e_parking/',
        help='path to save sensor data (default: ./e2e_parking/)')
    argparser.add_argument(
        '--task_num',
        default=16,
        type=int,
        help='number of parking task (default: 16')
    argparser.add_argument(
        '--map',
        default='Town04_Opt',
        help='map of carla (default: Town04_Opt)',
        choices=['Town04_Opt', 'Town05_Opt'])
    argparser.add_argument(
        '--random_seed',
        default=0,
        help='random seed to initialize env; if sets to 0, use current timestamp as seed (default: 0)')
    argparser.add_argument(
        '--bev_render_device',
        default='cpu',
        help='device used for BEV Rendering (default: cpu)',
        choices=['cpu', 'cuda'])
    argparser.add_argument(
        '--shuffle_veh',
        default=True,
        type=str2bool,
        help='shuffle static vehicles between tasks (default: True)')
    argparser.add_argument(
        '--shuffle_weather',
        default=False,
        type=str2bool,
        help='shuffle weather between tasks (default: False)')
    args = argparser.parse_args()
    args.width, args.height = [int(x) for x in args.res.split('x')]
    pygame.init()
    pygame.font.init()

    _world, ego_vehicle = world_init(args)
    # TODO 里面会查找vehicle 确定frame， 但是如果有多个vehicle 呢
    bev_view = generate_bev(args, _world, ego_vehicle)
    torch.save(bev_view, 'bev_view.pth')
    print(type(bev_view))
    print(bev_view.shape)
    img1 = encode_npy_to_pil(np.asarray(bev_view.squeeze().cpu()))
    print('bev_view saved')
    cv2.imwrite('bev_view.png', img1)


def view_bev():
    bev_view = torch.load('bev_view.pth')
    img1 = encode_npy_to_pil(np.asarray(bev_view.squeeze().cpu()))
    print(img1.shape)
    matrix = np.transpose(img1, (1, 2, 0))
    image = Image.fromarray(matrix)
    image.save('output_image.jpg')



if __name__ == '__main__':
    # main()
    view_bev()

